#include "gtest/gtest.h"
// #include "DataStructure/Common/include/Constant.h"
#include "sgPArpack.h"

// using SG::DataStructure::Common::TOLERANCE;
using namespace SG::Algebra;

class PArpackTest : public ::testing::Test
{
  public:
    static void SetUpTestSuite ()
    {}
};

/* test program to solve for the 9 smallest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 */
TEST_F (PArpackTest, eig3_d_lm)
{
    const int N   = 100;
    const int nev = 9;

    SpMat_t A (N, N);

    const int row_start = A.row_start ();
    const int row_end   = A.row_end ();

    for (int i = row_start; i < row_end; ++i)
    {
        A.set (i, i, N - i);
    }
    A.assemble ();

    SpMat_t B (N, N);
    for (int i = row_start; i < row_end; ++i)
    {
        B.set (i, i, 1);
    }
    B.assemble ();

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    PArpack<Real_t> solver;
    auto            opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "LM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Real_t> eigenvalues;
    std::vector<Vec_t>  eigenvectors;

    solver.solve (A, B, eigenvalues, eigenvectors);
    auto info = solver.info ();

    // Check
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<Real_t> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        exact_eigenvalues.push_back (i + 1);
    }

    // check the eigenvalues
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i], eigenvalues[i], TOLERANCE);
    }
}

/* test program to solve for the 9 largest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 */
TEST_F (PArpackTest, eig3_d_sm)
{
    const int N   = 100;
    const int nev = 9;

    SpMat_t A (N, N);

    const int row_start = A.row_start ();
    const int row_end   = A.row_end ();

    for (int i = row_start; i < row_end; ++i)
    {
        A.set (i, i, N - i);
    }
    A.assemble ();

    SpMat_t B (N, N);
    for (int i = row_start; i < row_end; ++i)
    {
        B.set (i, i, 1);
    }
    B.assemble ();

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    PArpack<Real_t> solver;
    auto            opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "SM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Real_t> eigenvalues;
    std::vector<Vec_t>  eigenvectors;
    solver.solve (A, B, eigenvalues, eigenvectors);

    // Check
    auto info = solver.info ();
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<double> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        exact_eigenvalues.push_back (N - nev + i + 1);
    }

    // check the eigenvalues
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i], eigenvalues[i], TOLERANCE);
    }
}

/* test program to solve for the 9 smallest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 *
 */
TEST_F (PArpackTest, eig3_z_lm)
{
    const int N   = 100;
    const int nev = 9;

    ComplexSpMat_t A (N, N);

    const int row_start = A.row_start ();
    const int row_end   = A.row_end ();

    for (int i = row_start; i < row_end; ++i)
    {
        Complex_t value (N - i, 0.0);
        A.set (i, i, value);
    }
    A.assemble ();

    ComplexSpMat_t B (N, N);

    for (int i = row_start; i < row_end; ++i)
    {
        Complex_t value (1.0, 0.0);
        B.set (i, i, value);
    }
    B.assemble ();

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    PArpack<Complex_t> solver;
    auto               opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "LM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Complex_t>    eigenvalues;
    std::vector<ComplexVec_t> eigenvectors;

    solver.solve (A, B, eigenvalues, eigenvectors);
    auto info = solver.info ();

    // Check
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<Complex_t> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        Complex_t value (i + 1, 0.0);
        exact_eigenvalues.push_back (i + 1);
    }

    // // check the eigenvalues
    // // For TArpack<double>, it seems that the converged Ritz values are returned in
    // // a way which is different from TArpack<double>.
    // // see https://github.com/opencollab/arpack-ng/blob/master/SRC/znaupd.f
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i].real (), eigenvalues[i].real (), TOLERANCE);
        ASSERT_NEAR (exact_eigenvalues[i].imag (), eigenvalues[i].imag (), TOLERANCE);
    }
}

/* test program to solve for the 9 largest eigenvalues of
 * A*x = lambda*x where A is the diagonal matrix
 * with entries 100, 99, ... , 2, 1 on the diagonal.
 */
TEST_F (PArpackTest, eig3_z_sm)
{
    const int N   = 100;
    const int nev = 9;

    ComplexSpMat_t A (N, N);

    const int      row_start = A.row_start ();
    const int      row_end   = A.row_end ();

    for (int i = row_start; i < row_end; ++i)
    {
        Complex_t value (N - i, 0.0);
        A.set (i, i, value);
    }
    
    A.assemble ();

    ComplexSpMat_t B (N, N);

    for (int i = row_start; i < row_end; ++i)
    {
        Complex_t value (1.0, 0.0);
        B.set (i, i, value);
    }
    B.assemble ();

    // solve the generalized eigenvalue problem A*x = lambda*B*x
    PArpack<Complex_t> solver;
    auto               opt = solver.getDefaultOption ();
    // SM for OP= (inv[K - sigma*M])*M, while LA for K*x = lambda*M*x
    opt.which = "SM";
    opt.nev   = 9;
    solver.setOption (opt);

    std::vector<Complex_t>    eigenvalues;
    std::vector<ComplexVec_t> eigenvectors;

    solver.solve (A, B, eigenvalues, eigenvectors);

    // Check
    auto info = solver.info ();
    ASSERT_TRUE (Info_t::Success == info);

    // the exact eigenvalues
    std::vector<Complex_t> exact_eigenvalues;
    for (int i = 0; i < nev; ++i)
    {
        Complex_t value (N - i);
        exact_eigenvalues.push_back (value);
    }

    // check the eigenvalues
    // For TArpack<Complex_t>, it seems that the converged Ritz values are returned in
    // a way which is different from TArpack<double>.
    // see https://github.com/opencollab/arpack-ng/blob/master/SRC/znaupd.f
    for (int i = 0; i < nev; ++i)
    {
        ASSERT_NEAR (exact_eigenvalues[i].real (), eigenvalues[i].real (), TOLERANCE);
        ASSERT_NEAR (exact_eigenvalues[i].imag (), eigenvalues[i].imag (), TOLERANCE);
    }
}